import random

# Pomocne funkcije
def mod_pow(a, n, m):
    result = 1
    a = a % m
    while n > 0:
        if n % 2 == 1:
            result = (result * a) % m
        a = (a * a) % m
        n = n // 2
        
    return result

def miller_rabin(n, k):
    if n <= 3:
        if n == 1:
            return False
        return True
  
    # n prost => n neparan => n = (2 ^ r) * d + 1
    d = n - 1
    r = 0
    while d % 2 == 0:
        r = r + 1
        d = d // 2
        
    for i in range(k):
        a = random.randrange(2, n - 1)
        
        x = mod_pow(a, d, n)

        if x == 1 or x == n - 1:
            continue
            
        wittness = True
        
        for j in range(r - 1):
            x = mod_pow(x, 2, n)
            if x == 1:
                return False
            if x == n - 1: # n - 1 = -1 (mod n)
                wittness = False
                break
        
        if wittness:
            return False
    return True


def get_prime(limit, k = 20):
    is_prime = False
    while not is_prime:
        n = random.randrange(limit)
        is_prime = miller_rabin(n, k)
    return n

# Pomoćna funkcija, prošireni Euklidov algoritam
def gcd(a, b):
    if b == 0:
        return a
    return gcd(b, a % b)

# Pomoćna funkcija, prošireni Euklidov algoritam
def ext_gcd(a, b):
    if b == 0:
        return (a, 1, 0)
    g, x, y = ext_gcd(b, a % b)
    return (g, y, x - a // b * y)

def mod_inv(a, m):
    g, x, y = ext_gcd(a, m)
    if g != 1:
        print("Vrednosti a i m nisu uzajamno proste!")
    else:
        return x % m

class RSA:
    def __init__(self, limit):
        self.limit = limit
        self.p = get_prime(2 ** (limit // 2))
        self.q = get_prime(2 ** (limit // 2))
        
        self.n = self.p * self.q
        
        self.phi = (self.p - 1) * (self.q - 1)
        
        self.generate_keys()
        
    def generate_keys(self):
        while True:
            e = random.randrange(2, self.phi-1)
            if gcd(e, self.phi) == 1:
                self.e = e
                break
                
        self.d = mod_inv(self.e, self.phi)
        
    def encrypt(self, m, e, n):
        return pow(m, e, n)
    
    def decrypt(self, me):
        return pow(me, self.d, self.n)

A = RSA(256)
B = RSA(256)

A_pub = (A.e, A.n)
B_pub = (B.e, B.n)


m1 = 123
print(f'A sent message 1: {m1}')

m1e = A.encrypt(m1, B_pub[0], B_pub[1])

print(f'A --[{m1e}] --> B')

m1d = B.decrypt(m1e)

print(f'B received message 1: {m1d}')

print()

m2 = 456
print(f'B sent message 2: {m2}')

m2e = B.encrypt(m2, A_pub[0], A_pub[1])

print(f'A <--[{m2e}] -- B')

m2d = A.decrypt(m2e)

print(f'A received message 2: {m2d}')

A sent message 1: 123
A --[51768325324231357949255474101733802782143864415793185897624686084288355934207] --> B
B received message 1: 123

B sent message 2: 456
A <--[8379748146270704912320313559995173099666599008890378715937326184404244640981] -- B
A received message 2: 456